import tf from '../tf'
import { ResidualDenseBlock } from './ResidualDenseBlock'

// 定义残差密集块 (RRDB)
export class RRDB extends tf.layers.Layer {
  // implements AddWeightAgent
  rdb1: ResidualDenseBlock
  rdb2: ResidualDenseBlock
  rdb3: ResidualDenseBlock
  constructor(numFeat: number, numGrowCh = 32) {
    super({})
    this.rdb1 = new ResidualDenseBlock(numFeat, numGrowCh)
    this.rdb2 = new ResidualDenseBlock(numFeat, numGrowCh)
    this.rdb3 = new ResidualDenseBlock(numFeat, numGrowCh)
  }

  call(x: any) {
    let out = this.rdb1.apply(x)
    out = this.rdb2.apply(out)
    out = this.rdb3.apply(out)
    if (x instanceof tf.Tensor) {
      out = tf.mul(out as tf.Tensor, 0.2)
    }
    return tf.layers.add().apply([out, x]) as any
  }

  build(inputShape: tf.Shape): void {
    const trainableWeights: tf.LayerVariable[] = []
    this.rdb1.build(inputShape)
    this.rdb2.build(inputShape)
    this.rdb3.build(inputShape)
    trainableWeights.push(...this.rdb1.trainableWeights)
    trainableWeights.push(...this.rdb2.trainableWeights)
    trainableWeights.push(...this.rdb3.trainableWeights)
    this._trainableWeights = trainableWeights
    this.built = true
  }
  setWeights(weights: tf.Tensor[]) {
    const rdb1Weights = weights.slice(0, 10)
    const rdb2Weights = weights.slice(10, 20)
    const rdb3Weights = weights.slice(20, 30)
    this.rdb1.setWeights(rdb1Weights)
    this.rdb2.setWeights(rdb2Weights)
    this.rdb3.setWeights(rdb3Weights)
  }
  getWeights(trainableOnly?: boolean): tf.Tensor[] {
    return [
      ...this.rdb1.getWeights(trainableOnly),
      ...this.rdb2.getWeights(trainableOnly),
      ...this.rdb3.getWeights(trainableOnly),
    ]
  }
  static get className() {
    return 'RRDB'
  }
}
tf.serialization.registerClass(RRDB)
// const model = new RRDB(64) // Example with 64 feature channels
// const input = tf.randomNormal([1, 64, 64, 64]) // Example input tensor
// const output = model.apply(input)
// @ts-ignore
// console.log(output.shape) // [1, 64, 64, 64];
